#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
标准气体车间三阶段调度（单文件版，含四算法+批量算例+图表导出）
========================================================
- 三阶段：配气(S1, 单机) -> 混匀(S2, 单机) -> 分析(S3, 多资源并行串联)
- 目标：最小化最大完工时间 C_max
- 算法：遗传算法(GA) / 模拟退火(SA) / 粒子群(PSO, Random Keys) / 禁忌搜索(TS)
- 功能：
  1) 生成 24 个由小到大的算例，保存到 `instances/`（JSON，可读性强）
  2) 若已存在算例，则直接加载，让用户选择算例编号
  3) 对所选算例运行四算法，保存：结果 JSON、**甘特图**、**收敛曲线**、**四算法对比收敛图** 到 `outputs/<算例名>/`

运行示例：
----------
python gas_shop_all_in_one.py --regen         # 重新生成 24 个算例
python gas_shop_all_in_one.py --id 5          # 直接选择第 5 个算例运行四算法
python gas_shop_all_in_one.py                 # 若未提供 --id，会交互式选择
"""

# 注：本文件兼容 Python 3.10+（推荐 3.11/3.12）。

import argparse
import json
import math
import os
from pathlib import Path
import random
import sys
import time
from dataclasses import dataclass
from typing import Dict, List, Tuple, Sequence

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

# —— 为中文字符设置字体，避免 DejaVu Sans 缺字警告 ——
matplotlib.rcParams['font.sans-serif'] = [
    'Microsoft YaHei', 'SimHei', 'Noto Sans CJK SC', 'Arial Unicode MS', 'DejaVu Sans'
]
matplotlib.rcParams['axes.unicode_minus'] = False

# =====================
# 基本数据结构
# =====================
Component = str  # 组分类别，如 'A','B','C',...

@dataclass
class Job:
    id: int
    components: List[Component]
    p_fill: int
    p_mix: int
    p_ana: Dict[Component, int]  # 每组分分析时间

    def to_dict(self) -> dict:
        return {
            'id': int(self.id),
            'components': list(self.components),
            'p_fill': int(self.p_fill),
            'p_mix': int(self.p_mix),
            'p_ana': {str(k): int(v) for k, v in self.p_ana.items()},
        }

    @staticmethod
    def from_dict(d: dict) -> 'Job':
        return Job(id=int(d['id']), components=list(d['components']),
                   p_fill=int(d['p_fill']), p_mix=int(d['p_mix']),
                   p_ana={str(k): int(v) for k, v in d['p_ana'].items()})

@dataclass
class Instance:
    jobs: List[Job]
    components: List[Component]
    analyzers_per_component: Dict[Component, int]
    meta: Dict[str, int]  # 记录生成参数

    def to_dict(self) -> dict:
        return {
            'jobs': [j.to_dict() for j in self.jobs],
            'components': list(self.components),
            'analyzers_per_component': {str(k): int(v) for k, v in self.analyzers_per_component.items()},
            'meta': {str(k): int(v) for k, v in self.meta.items()},
        }

    @staticmethod
    def from_dict(d: dict) -> 'Instance':
        return Instance(
            jobs=[Job.from_dict(x) for x in d['jobs']],
            components=[str(x) for x in d['components']],
            analyzers_per_component={str(k): int(v) for k, v in d['analyzers_per_component'].items()},
            meta={str(k): int(v) for k, v in d.get('meta', {}).items()},
        )

@dataclass
class ScheduleResult:
    order: List[int]
    cmax: int
    times_s1: Dict[int, Tuple[int, int]]
    times_s2: Dict[int, Tuple[int, int]]
    times_s3: Dict[int, Tuple[int, int]]

# =====================
# 工具函数（含 Numpy 类型安全的 JSON 导出）
# =====================

def ensure_dir(path: Path) -> None:
    path.mkdir(parents=True, exist_ok=True)


def _json_default(o):
    """把 numpy 类型安全地转换为原生 Python 类型，便于 json.dump。"""
    if isinstance(o, (np.integer,)):
        return int(o)
    if isinstance(o, (np.floating,)):
        return float(o)
    if isinstance(o, (np.ndarray,)):
        return o.tolist()
    # 也可在此扩展对 Path、set 等的处理
    raise TypeError(f"Object of type {o.__class__.__name__} is not JSON serializable")


def save_json(obj: dict, path: Path) -> None:
    ensure_dir(path.parent)
    with path.open('w', encoding='utf-8') as f:
        json.dump(obj, f, ensure_ascii=False, indent=2, default=_json_default)


def load_json(path: Path) -> dict:
    with path.open('r', encoding='utf-8') as f:
        return json.load(f)

# =====================
# 算例生成（24 个由小到大）
# =====================

def gen_instance(seed: int, n_jobs: int, n_types: int,
                 fill_base: int = 20, fill_per: int = 9,
                 mix_base: int = 0, mix_per: int = 7,
                 machines_per_type: int = 2) -> Instance:
    """按参数生成一个实例。"""
    rng = random.Random(seed)
    components = [chr(ord('A') + i) for i in range(n_types)]

    jobs: List[Job] = []
    for j in range(n_jobs):
        k = rng.randint(2, min(5, n_types))  # 每瓶组分数
        comps = rng.sample(components, k)
        p_fill = int(fill_base + fill_per * k + rng.randint(0, 5))
        p_mix = int(mix_base + mix_per * k + rng.randint(0, 4))
        p_ana: Dict[Component, int] = {}
        for c in comps:
            base = 12 + 3 * components.index(c)
            p_ana[c] = int(base + rng.randint(0, 10))
        jobs.append(Job(id=j, components=comps, p_fill=p_fill, p_mix=p_mix, p_ana=p_ana))

    analyzers = {c: machines_per_type for c in components}
    meta = dict(seed=int(seed), n_jobs=int(n_jobs), n_types=int(n_types),
                fill_base=int(fill_base), fill_per=int(fill_per),
                mix_base=int(mix_base), mix_per=int(mix_per),
                machines_per_type=int(machines_per_type))

    return Instance(jobs=jobs, components=components,
                    analyzers_per_component=analyzers, meta=meta)


def build_24_instances(inst_dir: Path, base_seed: int = 2025, regen: bool = False) -> List[Path]:
    """生成 24 个由小到大的算例，保存 JSON；返回文件路径列表。"""
    ensure_dir(inst_dir)

    files = sorted(inst_dir.glob('inst_*.json'))
    if not regen and len(files) >= 24:
        return files[:24]

    # 清理旧文件
    for p in inst_dir.glob('inst_*.json'):
        try:
            p.unlink()
        except Exception:
            pass

    specs: List[Tuple[int, int, int]] = []  # (n_jobs, n_types, machines_per_type)
    # 1~8: 小规模（10~24 瓶，4~5 类，设备各 2 台）
    for i in range(8):
        n_jobs = 10 + 2 * i
        n_types = 4 + (i // 4)
        mpt = 2
        specs.append((n_jobs, n_types, mpt))
    # 9~16: 中规模（26~40 瓶，6~7 类，设备各 3 台）
    for i in range(8):
        n_jobs = 26 + 2 * i
        n_types = 6 + (i // 4)
        mpt = 3
        specs.append((n_jobs, n_types, mpt))
    # 17~24: 大规模（44~72 瓶，8~9 类，设备各 4 台）
    for i in range(8):
        n_jobs = 44 + 4 * i
        n_types = 8 + (i // 6)
        mpt = 4
        specs.append((n_jobs, n_types, mpt))

    paths: List[Path] = []
    for idx, (n_jobs, n_types, mpt) in enumerate(specs, start=1):
        seed = base_seed + idx
        inst = gen_instance(seed=seed, n_jobs=n_jobs, n_types=n_types, machines_per_type=mpt)
        path = inst_dir / f"inst_{idx:02d}_J{n_jobs}_T{n_types}_M{mpt}.json"
        save_json(inst.to_dict(), path)
        paths.append(path)

    return paths

# =====================
# 调度评估（给定顺序）
# =====================

def evaluate_permutation(order: Sequence[int], ins: Instance) -> 'ScheduleResult':
    """S1/S2 单机队列；S3 多资源并行串联（同步占用，时长=各组分分析时间最大值）。"""
    # S1: 单机
    t1 = 0
    times_s1: Dict[int, Tuple[int, int]] = {}
    for j in order:
        s = t1
        e = s + int(ins.jobs[int(j)].p_fill)
        times_s1[int(j)] = (int(s), int(e))
        t1 = e

    # S2: 单机
    t2 = 0
    times_s2: Dict[int, Tuple[int, int]] = {}
    for j in order:
        j = int(j)
        ready = times_s1[j][1]
        s = max(int(ready), int(t2))
        e = s + int(ins.jobs[j].p_mix)
        times_s2[j] = (s, e)
        t2 = e

    # S3: 多资源并行串联
    avail: Dict[Component, List[int]] = {c: [0]*int(ins.analyzers_per_component[c]) for c in ins.components}
    times_s3: Dict[int, Tuple[int, int]] = {}

    for j in order:
        j = int(j)
        job = ins.jobs[j]
        ready = int(times_s2[j][1])
        if not job.components:
            times_s3[j] = (ready, ready)
            continue
        dur = max(int(job.p_ana[c]) for c in job.components)

        probe = ready
        while True:
            earliest = 0
            for c in job.components:
                machines = sorted(avail[c])
                earliest = max(earliest, machines[0])
            t0 = max(int(probe), int(earliest))
            feasible = True
            for c in job.components:
                if min(avail[c]) > t0:
                    feasible = False
                    break
            if feasible:
                for c in job.components:
                    machines = sorted(avail[c])
                    machines[0] = t0 + dur
                    avail[c] = machines
                times_s3[j] = (t0, t0 + dur)
                break
            else:
                probe = t0

    cmax = max((times_s3[j][1] for j in order), default=0)
    return ScheduleResult(order=[int(x) for x in order], cmax=int(cmax),
                          times_s1=times_s1, times_s2=times_s2, times_s3=times_s3)

# =====================
# 可视化：甘特图与收敛曲线
# =====================

def _color_for_job(jid: int):
    cmap = plt.get_cmap('tab20')
    return cmap(int(jid) % 20)


def plot_gantt(res: ScheduleResult, ins: Instance, out_path: Path, title: str = '') -> None:
    ensure_dir(out_path.parent)
    fig, ax = plt.subplots(figsize=(12, 4))
    lanes = ['S1:配气', 'S2:混匀', 'S3:分析(并行串联)']
    ybase = np.arange(len(lanes))

    for idx, times in enumerate([res.times_s1, res.times_s2, res.times_s3]):
        y = ybase[idx]
        for j in res.order:
            s, e = times[j]
            if e <= s:
                continue
            ax.add_patch(Rectangle((s, y-0.4), e-s, 0.8, facecolor=_color_for_job(j), edgecolor='black', linewidth=0.5))
            ax.text(s + (e - s)/2, y, f"J{j}", ha='center', va='center', fontsize=8)

    ax.set_yticks(ybase)
    ax.set_yticklabels(['S1:配气', 'S2:混匀', 'S3:分析(并行串联)'])
    ax.set_xlabel('时间 (min)')
    ax.set_title(title if title else f"甘特图  |  Cmax={res.cmax}")
    ax.set_ylim(-0.8, len(lanes)-0.2)
    ax.grid(True, axis='x', linestyle='--', alpha=0.3)
    fig.tight_layout()
    fig.savefig(out_path, dpi=160)
    plt.close(fig)

ess_line_styles = {'GA': '-', 'SA': '--', 'PSO': '-.', 'TS': ':'}


def plot_convergence(series: Dict[str, List[float]], out_path: Path, title: str = '收敛曲线') -> None:
    ensure_dir(out_path.parent)
    fig, ax = plt.subplots(figsize=(8, 4))
    for name, ys in series.items():
        xs = list(range(1, len(ys) + 1))
        style = ess_line_styles.get(name, '-')
        ax.plot(xs, [float(y) for y in ys], linestyle=style, label=name)
    ax.set_xlabel('迭代 / 代数')
    ax.set_ylabel('当前最优 Cmax')
    ax.set_title(title)
    ax.grid(True, linestyle='--', alpha=0.3)
    ax.legend()
    fig.tight_layout()
    fig.savefig(out_path, dpi=160)
    plt.close(fig)

# =====================
# 算法基类与四算法实现
# =====================

class Algorithm:
    name: str = 'ALG'
    def solve(self, ins: Instance):
        raise NotImplementedError

# —— 通用工具 ——

def random_permutation(n: int, rng: random.Random) -> List[int]:
    arr = list(range(n))
    rng.shuffle(arr)
    return arr


def decode_random_keys(keys: np.ndarray) -> List[int]:
    # 返回 Python int 而非 numpy.int64，避免 JSON 序列化报错
    return list(map(int, np.argsort(keys, kind='mergesort')))

# —— 遗传算法 ——
class GAAlgorithm(Algorithm):
    name = 'GA'
    def __init__(self, seed=0, pop_size=60, gens=200, cx_rate=0.9, mut_rate=0.2, tournament_k=3):
        self.seed = seed
        self.pop_size = pop_size
        self.gens = gens
        self.cx_rate = cx_rate
        self.mut_rate = mut_rate
        self.tournament_k = tournament_k

    @staticmethod
    def _ox(p1: List[int], p2: List[int], rng: random.Random) -> List[int]:
        n = len(p1)
        a, b = sorted(rng.sample(range(n), 2))
        child = [-1] * n
        child[a:b+1] = p1[a:b+1]
        fill = [x for x in p2 if x not in child]
        idx = 0
        for i in range(n):
            if child[i] == -1:
                child[i] = fill[idx]
                idx += 1
        return child

    def solve(self, ins: Instance):
        rng = random.Random(self.seed)
        n = len(ins.jobs)
        pop = [random_permutation(n, rng) for _ in range(self.pop_size)]
        fits = [evaluate_permutation(p, ins).cmax for p in pop]
        history: List[float] = [min(fits)]

        for _ in range(self.gens):
            new_pop: List[List[int]] = []
            elite_idx = int(np.argmin(fits))
            new_pop.append(pop[elite_idx])

            def select_one() -> List[int]:
                cand = rng.sample(range(len(pop)), self.tournament_k)
                i = min(cand, key=lambda idx: fits[idx])
                return pop[i]

            while len(new_pop) < self.pop_size:
                p1, p2 = select_one(), select_one()
                if rng.random() < self.cx_rate:
                    c = self._ox(p1, p2, rng)
                else:
                    c = p1.copy()
                if rng.random() < self.mut_rate:
                    i, j = rng.sample(range(n), 2)
                    c[i], c[j] = c[j], c[i]
                new_pop.append(c)

            pop = new_pop[:self.pop_size]
            fits = [evaluate_permutation(p, ins).cmax for p in pop]
            history.append(min(fits))

        best_idx = int(np.argmin(fits))
        best_p = pop[best_idx]
        return evaluate_permutation(best_p, ins), history

# —— 模拟退火 ——
class SAAlgorithm(Algorithm):
    name = 'SA'
    def __init__(self, seed=0, iters=20000, t_init=1000.0, t_final=1e-3, cooling=0.995):
        self.seed = seed
        self.iters = iters
        self.t_init = t_init
        self.t_final = t_final
        self.cooling = cooling

    def solve(self, ins: Instance):
        rng = random.Random(self.seed)
        n = len(ins.jobs)
        cur = random_permutation(n, rng)
        cur_fit = evaluate_permutation(cur, ins).cmax
        best, best_fit = cur, cur_fit
        T = self.t_init
        history = [best_fit]

        step = 0
        while T > self.t_final and step < self.iters:
            step += 1
            i, j = rng.sample(range(n), 2)
            nxt = cur.copy()
            nxt[i], nxt[j] = nxt[j], nxt[i]
            nxt_fit = evaluate_permutation(nxt, ins).cmax
            delta = nxt_fit - cur_fit
            if delta < 0 or rng.random() < math.exp(-delta / max(1e-9, T)):
                cur, cur_fit = nxt, nxt_fit
                if cur_fit < best_fit:
                    best, best_fit = cur, cur_fit
            history.append(best_fit)
            T *= self.cooling

        return evaluate_permutation(best, ins), history

# —— 粒子群（Random Keys） ——
class PSOAlgorithm(Algorithm):
    name = 'PSO'
    def __init__(self, seed=0, swarm_size=40, iters=400, w=0.7, c1=1.5, c2=1.5):
        self.seed = seed
        self.swarm_size = swarm_size
        self.iters = iters
        self.w = w
        self.c1 = c1
        self.c2 = c2

    def solve(self, ins: Instance):
        rng = np.random.default_rng(self.seed)
        n = len(ins.jobs)
        X = rng.uniform(0.0, 1.0, size=(self.swarm_size, n))
        V = rng.normal(0.0, 0.1, size=(self.swarm_size, n))

        def fitness_row(xrow: np.ndarray) -> int:
            order = decode_random_keys(xrow)
            return int(evaluate_permutation(order, ins).cmax)

        fits = np.array([fitness_row(X[i]) for i in range(self.swarm_size)])
        pbest = X.copy()
        pbest_fit = fits.copy()
        g_idx = int(np.argmin(fits))
        gbest = X[g_idx].copy()
        gbest_fit = float(fits[g_idx])
        history = [gbest_fit]

        for _ in range(self.iters):
            r1 = rng.uniform(0.0, 1.0, size=(self.swarm_size, n))
            r2 = rng.uniform(0.0, 1.0, size=(self.swarm_size, n))
            V = self.w * V + self.c1 * r1 * (pbest - X) + self.c2 * r2 * (gbest - X)
            X = np.clip(X + V, 0.0, 1.0)

            fits = np.array([fitness_row(X[i]) for i in range(self.swarm_size)])
            improved = fits < pbest_fit
            pbest[improved] = X[improved]
            pbest_fit[improved] = fits[improved]

            g_idx = int(np.argmin(pbest_fit))
            if float(pbest_fit[g_idx]) < gbest_fit:
                gbest = pbest[g_idx].copy()
                gbest_fit = float(pbest_fit[g_idx])
            history.append(gbest_fit)

        best_order = decode_random_keys(gbest)
        return evaluate_permutation(best_order, ins), history

# —— 禁忌搜索 ——
class TSAlgorithm(Algorithm):
    name = 'TS'
    def __init__(self, seed=0, iters=5000, tenure=20, cand_per_iter=50):
        self.seed = seed
        self.iters = iters
        self.tenure = tenure
        self.cand_per_iter = cand_per_iter

    def solve(self, ins: Instance):
        rng = random.Random(self.seed)
        n = len(ins.jobs)
        cur = random_permutation(n, rng)
        cur_fit = evaluate_permutation(cur, ins).cmax
        best, best_fit = cur, cur_fit
        history = [best_fit]

        tabu: Dict[Tuple[int, int], int] = {}
        for _ in range(self.iters):
            candidates: List[Tuple[int, int, List[int], int]] = []
            tried = set()
            while len(candidates) < self.cand_per_iter:
                i, j = sorted(rng.sample(range(n), 2))
                if (i, j) in tried:
                    continue
                tried.add((i, j))
                q = cur.copy()
                q[i], q[j] = q[j], q[i]
                f = evaluate_permutation(q, ins).cmax
                candidates.append((i, j, q, f))

            selected = None
            selected_fit = math.inf
            for (i, j, q, f) in candidates:
                move = tuple(sorted((cur[i], cur[j])))
                is_tabu = move in tabu
                if (not is_tabu and f < selected_fit) or (is_tabu and f < best_fit):
                    selected = (i, j, q, f, move)
                    selected_fit = f

            i, j, q, f, move = selected  # type: ignore
            cur, cur_fit = q, f

            tabu[move] = self.tenure
            expired = [k for k, v in tabu.items() if v <= 1]
            tabu = {k: v-1 for k, v in tabu.items() if k not in expired}

            if cur_fit < best_fit:
                best, best_fit = cur, cur_fit
            history.append(best_fit)

        return evaluate_permutation(best, ins), history

# =====================
# 主流程
# =====================

def list_instances(inst_dir: Path) -> List[Path]:
    return sorted(inst_dir.glob('inst_*.json'))


def select_instance_interactive(files: List[Path]) -> int:
    print('可用算例：')
    for i, p in enumerate(files, start=1):
        print(f"[{i:02d}] {p.name}")
    while True:
        try:
            s = input('请输入算例编号 (1-{:02d}): '.format(len(files))).strip()
            idx = int(s)
            if 1 <= idx <= len(files):
                return idx
        except Exception:
            pass
        print('输入无效，请重试。')


def run_all_algorithms(ins: Instance, out_dir: Path, base_seed: int = 2025) -> None:
    ensure_dir(out_dir)

    algos: List[Algorithm] = [
        GAAlgorithm(seed=base_seed+1, pop_size=60, gens=200, cx_rate=0.9, mut_rate=0.2, tournament_k=3),
        SAAlgorithm(seed=base_seed+2, iters=20000, t_init=1000.0, t_final=1e-3, cooling=0.995),
        PSOAlgorithm(seed=base_seed+3, swarm_size=40, iters=400, w=0.7, c1=1.5, c2=1.5),
        TSAlgorithm(seed=base_seed+4, iters=5000, tenure=20, cand_per_iter=50),
    ]

    combined_series: Dict[str, List[float]] = {}

    for alg in algos:
        print(f"===== 运行 {alg.name} =====")
        t0 = time.time()
        res, hist = alg.solve(ins)
        t1 = time.time()
        print(f"{alg.name} 完成：Cmax = {res.cmax}  |  耗时 = {t1 - t0:.2f}s")

        # 结果保存（确保所有 numpy 类型被转为原生类型）
        save_json({
            'algo': alg.name,
            'cmax': int(res.cmax),
            'order': [int(x) for x in res.order],
            'times_s1': {str(k): [int(v[0]), int(v[1])] for k, v in res.times_s1.items()},
            'times_s2': {str(k): [int(v[0]), int(v[1])] for k, v in res.times_s2.items()},
            'times_s3': {str(k): [int(v[0]), int(v[1])] for k, v in res.times_s3.items()},
            'solve_time_sec': float(t1 - t0),
        }, out_dir / f"{alg.name}_result.json")

        # 甘特图
        plot_gantt(res, ins, out_dir / f"{alg.name}_gantt.png", title=f"{alg.name} 甘特图  |  Cmax={res.cmax}")

        # 收敛曲线（单算法）
        plot_convergence({alg.name: [float(x) for x in hist]}, out_dir / f"{alg.name}_convergence.png", title=f"{alg.name} 收敛曲线")
        combined_series[alg.name] = [float(x) for x in hist]

    plot_convergence(combined_series, out_dir / f"ALL_convergence.png", title="四算法收敛曲线对比")
    print(f"图片与结果已保存到: {out_dir.resolve()}")


def main():
    parser = argparse.ArgumentParser(description='标准气体车间三阶段调度：四算法对比（单文件版）')
    parser.add_argument('--regen', action='store_true', help='强制重新生成 24 个算例')
    parser.add_argument('--id', type=int, default=None, help='直接指定算例编号 (1-24)')
    parser.add_argument('--base_seed', type=int, default=2025, help='基础随机种子（影响算例与算法）')
    args = parser.parse_args()

    root = Path.cwd()
    inst_dir = root / 'instances'
    out_root = root / 'outputs'

    files = build_24_instances(inst_dir, base_seed=args.base_seed, regen=args.regen)
    print(f"共 {len(files)} 个算例位于: {inst_dir.resolve()}")

    if args.id is None:
        idx = select_instance_interactive(files)
    else:
        if not (1 <= args.id <= len(files)):
            print(f"指定的 --id 无效，应在 1~{len(files)} 之间。", file=sys.stderr)
            sys.exit(2)
        idx = args.id

    inst_path = files[idx-1]
    print(f"选择算例 [{idx:02d}]: {inst_path.name}")
    ins = Instance.from_dict(load_json(inst_path))

    out_dir = out_root / inst_path.stem
    run_all_algorithms(ins, out_dir=out_dir, base_seed=args.base_seed)


if __name__ == '__main__':
    main()
